import torch
from torch.utils.data import Dataset
import pandas as pd
from PIL import Image
from torchvision import transforms
import json
from utils.process_labels import encode_labels
# 定义自己的类
class BottleDataset(Dataset):

    # 初始化
    def __init__(self, csv_path,data_local,base_size=16,is_train=True):
        self.base_size = base_size
        self.is_train = is_train
        # 读入数据
        fh = open(csv_path, 'r',encoding='UTF-8')
        self.imgs = []
        img_dict={}
        for line in fh:
            line = line.rstrip()
            dict_json = json.loads(line)

            self.imgs.append((dict_json['source'].replace("F:/GitRepository/BottleClassification/BottleDataset/",data_local),dict_json['annotation'][0]['name']))
            img_dict[dict_json['annotation'][0]['name']]= img_dict.get(dict_json['annotation'][0]['name'],0)+1
        print("---img label count:",img_dict)
    # 返回df的长度
    def __len__(self):
        return len(self.imgs)

    # 获取第idx+1列的数据
    def __getitem__(self, index):
        fn, label_name = self.imgs[index]
        input_img = Image.open(fn).convert('RGB')
        label= encode_labels(label_name)
        if self.is_train:
            transformed_result = self.transform_tr(input_img)
        else:
            transformed_result = self.transform_val(input_img)
        # label =  transformed_result['label'].numpy()
        return (transformed_result,label)
    def transform_tr(self, input_img):
        composed_transforms = transforms.Compose([
            # tr.RandomHorizontalFlip(),
            # # tr.RandomScaleCrop(base_size=self.base_size, crop_size=self.base_size),
            # tr.RandomGaussianBlur(),
            # tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            # tr.ToTensor()
            transforms.RandomHorizontalFlip(p=0.8),
            transforms.RandomRotation(45),
            transforms.Resize(self.base_size),  # smaller edge of the image will be matched to this number
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # nomalilzation
        ])

        return composed_transforms(input_img)

    def transform_val(self, input_img):

        composed_transforms = transforms.Compose([
            # tr.FixScaleCrop(crop_size=513),
            # tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            # tr.ToTensor()

            transforms.Resize(self.base_size),  # smaller edge of the image will be matched to this number
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # nomalilzation
        ])

        return composed_transforms(input_img)
if __name__ == "__main__":
    ds = BottleDataset('train.csv')
    a= ds.__getitem__(5)